package top.v5it.japi.plus.core.spring.web;

import cn.hutool.core.map.MapUtil;
import cn.hutool.core.util.StrUtil;
import lombok.extern.slf4j.Slf4j;

import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.stream.Stream;

/**
 * xss过虑
 *
 * @author zhanpu
 * @date 20-12-24
 */
@Slf4j
public class JapiXssHttpServletRequestWrapper extends JapiHttpServletRequestWrapper {

    /**
     * 构造函数
     *
     * @param request           {@link HttpServletRequest}
     * @param isIncludeRichText true-过虑，false-不过虑
     */
    public JapiXssHttpServletRequestWrapper(HttpServletRequest request, boolean isIncludeRichText) throws IOException {
        super(request, isIncludeRichText);
    }

    /**
     * 覆盖getParameter方法，将参数名和参数值都做xss过滤。<br/>
     * 如果需要获得原始的值，则通过super.getParameterValues(name)来获取<br/>
     * getParameterNames,getParameterValues和getParameterMap也可能需要覆盖
     */
    @Override
    public String getParameter(String name) {
        Boolean flag = ("content".equals(name) || name.endsWith("WithHtml"));
        return Boolean.TRUE.equals(flag) && !isIncludeRichText
                ? super.getParameter(name)
                : super.clean(super.getParameter(super.clean(name)));
    }

    /**
     * 覆盖getParameterValues方法
     *
     * @param name
     * @return
     */
    @Override
    public String[] getParameterValues(String name) {
        return Stream
                .of(super.getParameterValues(super.clean(name)))
                .map(super::clean)
                .toArray(String[]::new);
    }

    /**
     * 覆盖getHeader方法，将参数名和参数值都做xss过滤。<br/>
     * 如果需要获得原始的值，则通过super.getHeaders(name)来获取<br/>
     * getHeaderNames 也可能需要覆盖
     */
    @Override
    public String getHeader(String name) {
        return super.clean(super.getHeader(super.clean(name)));
    }

    @Override
    public String getQueryString() {
        return super.clean(super.getQueryString());
    }

    @Override
    public Map<String, String[]> getParameterMap() {
        Map<String, String[]> parameterMap = super.getParameterMap();

        if (MapUtil.isEmpty(parameterMap)) {
            return parameterMap;
        }

        Map<String, String[]> rtvMap = new HashMap<>(parameterMap.size());
        parameterMap.forEach((k, v) ->
                rtvMap.put(super.clean(k)
                        , Stream.of(v)
                                .map(super::clean)
                                .toArray(String[]::new)));

        return rtvMap;
    }

    /**
     * 重写getInputStream
     *
     * @return
     * @throws IOException
     */
    @Override
    public ServletInputStream getInputStream() throws IOException {
        String tempBody = super.getBody();
        return StrUtil.isBlank(tempBody)
                ? super.getInputStream()
                : new JapiServletInputStream(new ByteArrayInputStream(super.clean(tempBody).getBytes()));
    }

}
